{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Example: Reading and writing an equilibrium to IMAS (via IMAS-Python)\n",
    "\n",
    "In this example, we take a MAST-U-like equilibrium and show how to save it to, and load it back from, the [IMAS](https://imas.iter.org/) `equilibrium` IDS using the [IMAS-Python](https://github.com/iterorganization/IMAS-Python) package. IDSs (Interface Data Structures) are the standardised data structures IMAS uses to link different modelling codes together; here we serialise the IDS to a netCDF file."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As in the previous example, we run our standard equilibrium solve - making it lower single null by pushing the P6 current."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "# build machine\n",
    "from freegsnke import build_machine\n",
    "tokamak = build_machine.tokamak(\n",
    "    active_coils_path=f\"../machine_configs/MAST-U/MAST-U_like_active_coils.pickle\",\n",
    "    passive_coils_path=f\"../machine_configs/MAST-U/MAST-U_like_passive_coils.pickle\",\n",
    "    limiter_path=f\"../machine_configs/MAST-U/MAST-U_like_limiter.pickle\",\n",
    "    wall_path=f\"../machine_configs/MAST-U/MAST-U_like_wall.pickle\",\n",
    ")\n",
    "\n",
    "# initialise equilibrium object\n",
    "from freegsnke import equilibrium_update\n",
    "eq = equilibrium_update.Equilibrium(\n",
    "    tokamak=tokamak,\n",
    "    Rmin=0.1, Rmax=2.0,   # radial range\n",
    "    Zmin=-2.2, Zmax=2.2,  # vertical range\n",
    "    nx=65,                # number of grid points in the radial direction (needs to be of the form (2**n + 1) with n being an integer)\n",
    "    ny=129,               # number of grid points in the vertical direction (needs to be of the form (2**n + 1) with n being an integer)\n",
    ")\n",
    "\n",
    "# initialise profile object\n",
    "from freegsnke.jtor_update import ConstrainPaxisIp\n",
    "profiles = ConstrainPaxisIp(\n",
    "    eq=eq,\n",
    "    paxis=8.1e3,\n",
    "    Ip=6.2e5,\n",
    "    fvac=0.5,\n",
    "    alpha_m=1.8,\n",
    "    alpha_n=1.2\n",
    ")\n",
    "\n",
    "# initialise solver\n",
    "from freegsnke import GSstaticsolver\n",
    "GSStaticSolver = GSstaticsolver.NKGSsolver(eq)\n",
    "\n",
    "# set coil currents\n",
    "import pickle\n",
    "with open('data/simple_diverted_currents_PaxisIp.pk', 'rb') as f:\n",
    "    current_values = pickle.load(f)\n",
    "\n",
    "for key in current_values.keys():\n",
    "    eq.tokamak[key].current = current_values[key]\n",
    "\n",
    "# change this to shift plasma vertically up (-) or down (+)\n",
    "eq.tokamak[\"P6\"].current += 500\n",
    "\n",
    "# carry out forward solve\n",
    "GSStaticSolver.solve(eq=eq,\n",
    "                     profiles=profiles,\n",
    "                     constrain=None,\n",
    "                     target_relative_tolerance=1e-9)\n",
    "\n",
    "# plot the resulting equilibrium\n",
    "fig1, ax1 = plt.subplots(1, 1, figsize=(4, 8), dpi=80)\n",
    "ax1.grid(True, which='both')\n",
    "eq.plot(axis=ax1, show=False)\n",
    "eq.tokamak.plot(axis=ax1, show=False)\n",
    "ax1.set_xlim(0.1, 2.15)\n",
    "ax1.set_ylim(-2.25, 2.25)\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Writing to the equilibrium IDS\n",
    "\n",
    "FreeGSNKE provides a helper module, `freegsnke.imas_read_write`, that fills in the relevant fields of an IMAS `equilibrium` IDS from a solved `eq`/`profiles` pair: high-level `ids_properties`/`code` metadata, `vacuum_toroidal_field`, and a single `time_slice` containing the plasma boundary, global quantities, 1D flux-surface profiles, and three 2D poloidal flux maps on FreeGSNKE's rectangular (R, Z) grid - `total` (also carrying `j_phi`), `plasma`-only, and `vacuum` (coil-only) - tagged using the standard `equilibrium_profiles_2d_identifier` codes.\n",
    "\n",
    "See the [equilibrium IDS data dictionary](https://imas-data-dictionary.readthedocs.io/en/latest/generated/ids/equilibrium.html) for the full structure."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from freegsnke import imas_read_write\n",
    "\n",
    "equilibrium_ids = imas_read_write.write_equilibrium_to_ids(eq, profiles)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that it has been built, we can inspect the overall structure using IMAS-Python's own utilities."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import imas\n",
    "\n",
    "# inspect at a high level\n",
    "imas.util.inspect(equilibrium_ids, hide_empty_nodes=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can inspect at lower levels of the structure too."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "imas.util.inspect(equilibrium_ids.time_slice[0].global_quantities, hide_empty_nodes=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If unsure of where exactly a quantity is located in the structure, you can search for certain strings. Note that this returns all possible locations in the data dictionary, without revealing whether that location is actually filled."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "imas.util.find_paths(equilibrium_ids, \"psi_boundary\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Writing to netCDF\n",
    "\n",
    "IMAS-Python IDSs are written to disk via an `imas.DBEntry`. `imas_read_write.save_equilibrium_ids` wraps this for the netCDF backend - see [here](https://imas-python.readthedocs.io/en/stable/netcdf.html) for more details."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "imas_read_write.save_equilibrium_ids(equilibrium_ids, \"MASTU_LSN_equilibrium.nc\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Reading from netCDF\n",
    "\n",
    "Now we can read the equilibrium IDS back from the netCDF file.\n",
    "\n",
    "An `equilibrium` IDS doesn't store individual coil currents (only the boundary, global quantities, and profile data), so - just as with the geqdsk example - we use the inverse solver to re-build the equilibrium from what's actually saved.\n",
    "\n",
    "We build the constraints entirely from what the IDS gives us:\n",
    "\n",
    "- an **isoflux constraint** from the plasma boundary outline (`time_slice.boundary.outline`) - all these points should sit on the same flux contour, without us needing to know its absolute value;\n",
    "- **null-point constraints** at the magnetic axis (`time_slice.global_quantities.magnetic_axis`);\n",
    "- a **`psi_vals` constraint built from the *entire* saved 2D `psi` map** (`time_slice.profiles_2d[0].psi`) - unlike the constraints above, this one is given the actual flux *values*, not just which points should share a value, so it's what lets the optimiser pin down the absolute flux level (and, with it, the individual coil currents) rather than just the plasma shape. Passing the grid's own R/Z axes (rather than a flattened list of points) triggers `Inverse_optimizer`'s \"full grid\" fast path, which reuses the equilibrium's own cached Greens functions instead of recomputing them pointwise.\n",
    "\n",
    "To check that the IDS faithfully captures the equilibrium, we'll reconstruct it purely from what's in the file: the grid extent, the plasma current, and the `dpressure_dpsi`/`f_df_dpsi` profile data.\n",
    "\n",
    "FreeGSNKE's `GeneralPprimeFFprime` profile class accepts exactly this kind of tabulated profile data directly, rather than a parametric form like `ConstrainPaxisIp` above.\n",
    "\n",
    "We reset the coil currents on a **fresh, independent** tokamak to zero first, to make sure we aren't just reusing the answer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loaded_ids = imas_read_write.load_equilibrium_ids(\"MASTU_LSN_equilibrium.nc\")\n",
    "time_slice = loaded_ids.time_slice[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# grid data (IDS stores psi in Wb, i.e. with the 2*pi factor included - FreeGSNKE works in Wb/rad)\n",
    "R_read = np.array(time_slice.profiles_2d[0].grid.dim1)\n",
    "Z_read = np.array(time_slice.profiles_2d[0].grid.dim2)\n",
    "psi_read = np.array(time_slice.profiles_2d[0].psi) / (2 * np.pi)\n",
    "\n",
    "Rmin, Rmax = R_read[0], R_read[-1]\n",
    "Zmin, Zmax = Z_read[0], Z_read[-1]\n",
    "nx_read, ny_read = len(R_read), len(Z_read)\n",
    "\n",
    "# profile data\n",
    "psi_n_read = np.array(time_slice.profiles_1d.psi_norm)\n",
    "pprime_read = np.array(time_slice.profiles_1d.dpressure_dpsi)\n",
    "ffprime_read = np.array(time_slice.profiles_1d.f_df_dpsi)\n",
    "Ip_read = time_slice.global_quantities.ip\n",
    "fvac_read = loaded_ids.vacuum_toroidal_field.b0[0] * loaded_ids.vacuum_toroidal_field.r0\n",
    "\n",
    "# boundary and magnetic axis\n",
    "rbdry, zbdry = np.array(time_slice.boundary.outline.r), np.array(time_slice.boundary.outline.z)\n",
    "rmagx, zmagx = time_slice.global_quantities.magnetic_axis.r, time_slice.global_quantities.magnetic_axis.z"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's set up and run the inverse solver, using those quantities plus the constraints described above to reconstruct the coil currents."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from freegsnke.jtor_update import GeneralPprimeFFprime\n",
    "from freegsnke.inverse import Inverse_optimizer\n",
    "\n",
    "# coil names\n",
    "active_coil_names = tokamak.coils_list[: tokamak.n_active_coils]\n",
    "\n",
    "# fresh tokamak so we can set coil currents to zero and compare to old eq object\n",
    "tokamak_inv = build_machine.tokamak(\n",
    "    active_coils_path=f\"../machine_configs/MAST-U/MAST-U_like_active_coils.pickle\",\n",
    "    passive_coils_path=f\"../machine_configs/MAST-U/MAST-U_like_passive_coils.pickle\",\n",
    "    limiter_path=f\"../machine_configs/MAST-U/MAST-U_like_limiter.pickle\",\n",
    "    wall_path=f\"../machine_configs/MAST-U/MAST-U_like_wall.pickle\",\n",
    ")\n",
    "\n",
    "# fresh equilibrium object\n",
    "eq_inv = equilibrium_update.Equilibrium(\n",
    "    tokamak=tokamak_inv,\n",
    "    Rmin=Rmin, Rmax=Rmax,\n",
    "    Zmin=Zmin, Zmax=Zmax,\n",
    "    nx=nx_read,\n",
    "    ny=ny_read,\n",
    ")\n",
    "\n",
    "# fresh profiles object, built from the tabulated data saved to the IDS\n",
    "profiles_inv = GeneralPprimeFFprime(\n",
    "    eq=eq_inv,\n",
    "    Ip=Ip_read,\n",
    "    fvac=fvac_read,\n",
    "    psi_n=psi_n_read,\n",
    "    pprime_data=pprime_read,\n",
    "    ffprime_data=ffprime_read,\n",
    ")\n",
    "\n",
    "# isoflux constraint: subsample the boundary outline stored in the IDS\n",
    "isoflux_set = [[rbdry[::5], zbdry[::5]]]\n",
    "\n",
    "# null-point constraint: magnetic axis (from the IDS)\n",
    "null_points = [\n",
    "    [rmagx],\n",
    "    [zmagx],\n",
    "]\n",
    "\n",
    "# psi_vals constraint built from the *entire* saved 2D psi map (full-grid fast path)\n",
    "psi_vals = [eq_inv.R, eq_inv.Z, psi_read]\n",
    "\n",
    "# build constraints\n",
    "constrain = Inverse_optimizer(\n",
    "    isoflux_set=isoflux_set,\n",
    "    null_points=null_points,\n",
    "    psi_vals=psi_vals,\n",
    ")\n",
    "\n",
    "# initialise solver object\n",
    "GSStaticSolver_inv = GSstaticsolver.NKGSsolver(eq_inv)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set coils to zero current\n",
    "for label in active_coil_names:\n",
    "    eq_inv.tokamak.set_coil_current(label, 0.0)\n",
    "\n",
    "# default plasma flux guess\n",
    "eq_inv.set_plasma_psi(eq_inv.create_psi_plasma_default())\n",
    "\n",
    "# solve\n",
    "GSStaticSolver_inv.solve(\n",
    "    eq=eq_inv,\n",
    "    profiles=profiles_inv,\n",
    "    constrain=constrain,\n",
    "    target_relative_tolerance=1e-4,\n",
    "    target_relative_psit_update=1e-3,\n",
    "    max_solving_iterations=50,\n",
    "    l2_reg=1e-9,\n",
    "    verbose=True,\n",
    "    Picard_handover=1e-5,\n",
    "    full_jacobian_handover=[2e-2,1e-2],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This converges properly and gives the correct total flux, coil currents close to their true values, and `psi_axis`/`psi_bndry` matching to several decimal places."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# print coil currents\n",
    "reconstructed_currents = {label: eq_inv.tokamak[label].current for label in active_coil_names}\n",
    "original_active_currents = {label: eq.tokamak[label].current for label in active_coil_names}\n",
    "\n",
    "print(\"Coil currents [A]:\")\n",
    "print(f\"{'coil':>10s}  {'original':>12s}  {'reconstructed':>14s}\")\n",
    "for label in active_coil_names:\n",
    "    print(f\"{label:>10s}  {original_active_currents[label]:12.1f}  {reconstructed_currents[label]:14.1f}\")\n",
    "\n",
    "# some small differences\n",
    "psi_diff_inv = np.abs(eq.psi() - eq_inv.psi())\n",
    "print(f\"\\nMagnetic axis:   original = {eq.magneticAxis()[0:2]}, reconstructed = {eq_inv.magneticAxis()[0:2]}\")\n",
    "print(f\"psi_axis:        original = {eq.psi_axis:.6g}, reconstructed = {eq_inv.psi_axis:.6g}\")\n",
    "print(f\"psi_boundary:    original = {eq.psi_bndry:.6g}, reconstructed = {eq_inv.psi_bndry:.6g}\")\n",
    "print(f\"Max |Δpsi| relative to max|psi|: {np.max(psi_diff_inv) / np.max(np.abs(eq.psi())):.3e}\")\n",
    "\n",
    "# plot the two equilibria side by side\n",
    "fig1, (ax1, ax2) = plt.subplots(1, 2, figsize=(8, 8), dpi=80)\n",
    "for ax, e, title in [(ax1, eq, \"Original\"), (ax2, eq_inv, \"Reconstruction from IDS\")]:\n",
    "    ax.grid(True, which=\"both\")\n",
    "    e.plot(axis=ax, show=False)\n",
    "    e.tokamak.plot(axis=ax, show=False)\n",
    "    ax.set_xlim(0.1, 2.15)\n",
    "    ax.set_ylim(-2.25, 2.25)\n",
    "    ax.set_title(title)\n",
    "    if ax == ax2:\n",
    "        constrain.plot(axis=ax2)\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the future, we'll work on adding in capabilities to save the machine description and coil currents to the `pf_active` and `pf_passive` IDS formats. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "freegsnke",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.20"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}